﻿#region License Info
//Component of Cronos Package, http://www.codeplex.com/cronos
//Copyright (C) 2009 Anthony Brockwell

//This program is free software; you can redistribute it and/or
//modify it under the terms of the GNU General Public License
//as published by the Free Software Foundation; either version 2
//of the License, or (at your option) any later version.

//This program is distributed in the hope that it will be useful,
//but WITHOUT ANY WARRANTY; without even the implied warranty of
//MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//GNU General Public License for more details.

//You should have received a copy of the GNU General Public License
//along with this program; if not, write to the Free Software
//Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#endregion


using System;
using System.Collections.Generic;
using System.Text;
using System.Windows.Forms;
using ABMath.IridiumExtensions;
using ABMath.ModelFramework.Data;
using ABMath.ModelFramework.Models;
using MathNet.Numerics.Distributions;
using MathNet.Numerics.LinearAlgebra;
using MathNet.Numerics.RandomSources;

namespace ExamplePlugin
{
    [Serializable]
    public class SwitchingARMAModel : UnivariateTimeSeriesModel, IRealTimePredictable, IMLEEstimable
    {
        private TimeSeries predictiveP1;
        private TimeSeries predictiveMean;
        private TimeSeries MRIndicator;
        private TimeSeries TFIndicator;

        public double Mu
        {
            get { return Parameters[0]; }
            set { Parameters[0] = value; }
        }

        public double Sigma
        {
            get { return Parameters[1]; }
            set { Parameters[1] = value; }
        }

        public double Phi1
        {
            get { return Parameters[2]; }
            set { Parameters[2] = value; }
        }

        public double Phi2
        {
            get { return Parameters[3]; }
            set { Parameters[3] = value; }
        }

        public double PTransition11
        {
            get { return Parameters[4]; }
            set { Parameters[4] = value; }
        }

        public double PTransition22
        {
            get { return Parameters[5]; }
            set { Parameters[5] = value; }
        }

        public override string GetParameterName(int index)
        {
            switch (index)
            {
                case 0:
                    return "Mu";
                case 1:
                    return "Sigma";
                case 2:
                    return "Phi(1)";
                case 3:
                    return "Phi(2)";
                case 4:
                    return "P_11";
                case 5:
                    return "P_22";
                default:
                    throw new ArgumentException("Invalid parameter index");
            }
        }

        public override string GetParameterDescription(int index)
        {
            return null;
        }

        public double MarginalP1
        {
            get
            {
                double p1 = PTransition11;
                double p2 = PTransition22;
                return (p2 - 1)/(p1 + p2 - 2);
            }
        }

        public double MarginalP2
        {
            get
            {
                return 1 - MarginalP1;
            }
        }

        public SwitchingARMAModel()
        {
            LocalInitializeParameters();
        }

        public override string Description
        {
            get
            {
                var sb = new StringBuilder(1024);
                sb.AppendLine("Switching AR Model:");
                sb.AppendFormat("X(t) = Phi(t) X(t-1) + Z(t)");
                sb.AppendLine();
                sb.AppendFormat("Phi(t) switches between {0:0.000} and {1:0.000}{2}", Phi1, Phi2, Environment.NewLine);
                sb.AppendLine("with transition probability matrix");
                sb.AppendFormat("[ [{0:0.000}  {1:0.000}]   {2}", PTransition11, 1.0-PTransition11, Environment.NewLine);
                sb.AppendFormat("  [{0:0.000}  {1:0.000}] ] {2}", 1.0-PTransition22, PTransition22, Environment.NewLine);
                sb.AppendLine();
                sb.AppendFormat("Z(t) ~ N(0,{0:0.0000}^2) (iid){1}", Sigma, Environment.NewLine);
                sb.AppendFormat("Mean     = {0:0.0000}{1}", Mu, Environment.NewLine);
                sb.AppendLine();
                sb.AppendFormat("Marginal P(State 1) = {0:0.0000}{1}", MarginalP1, Environment.NewLine);
                return sb.ToString();
            }
        }

        public override string GetShortDescription()
        {
            return string.Format("Switching AR");
        }

        public override Vector Parameters
        {
            get; set;
        }

        public Vector ParameterToCube(Vector param)
        {
            var cube = new Vector(6);

            cube[0] = Math.Exp(param[0]) / (1 + Math.Exp(param[0])); // real to [0,1]
            cube[1] = param[1] / (1 + param[1]);                     // real+ to [0,1]
            cube[2] = (param[2] + 1) / 2.0;                          // [-1,1] to [0,1]
            cube[3] = (param[3] - param[2])/(1 - param[2]);
            cube[4] = param[4];                                      // [0,1] to [0,1]
            cube[5] = param[5];                                      // [0,1] to [0,1]

            return cube;
        }

        public Vector CubeToParameter(Vector cube)
        {
            var param = new Vector(6);

            param[0] = Math.Log(cube[0] / (1 - cube[0]));
            param[1] = cube[1] / (1 - cube[1]);
            param[2] = cube[2]*2 - 1;
            param[3] = (1 - param[2])*cube[3] + param[2];   // param[3] >= param[2]
            param[4] = cube[4];
            param[5] = cube[5];

            return param;
        }

        protected override bool CheckParameterValidity(Vector param)
        {
            bool valid = true;
            if (param[1] <= 0)
                valid = false;
            if (Math.Abs(param[2]) > 1)
                valid = false;
            if (Math.Abs(param[3]) > 1)
                valid = false;
            if (param[4] >= 1 || param[4] <= 0)
                valid = false;
            if (param[5] >= 1 || param[5] <= 0)
                valid = false;
            return valid;
        }

        public void CarryOutPreMLEComputations()
        {
            // nothing to do here: if there are some aspects of likelihood computation that can be reused with different parameters and the same data,
            // we would do that here
        }

        private void ComputeResidualsAndOutputs()
        {
            var resids = new TimeSeries();
            predictiveP1 = new TimeSeries();
            predictiveMean = new TimeSeries();
            MRIndicator = new TimeSeries();
            TFIndicator = new TimeSeries();
            GoodnessOfFit = ComputeLogLikelihood(Parameters, resids, predictiveP1, predictiveMean,
                                                 MRIndicator, TFIndicator);
            Residuals = resids;
            return;
        }

        private Matrix BuildTransitionMatrix()
        {
            var Q = new Matrix(2, 2);
            Q[0, 0] = PTransition11;
            Q[0, 1] = 1.0 - PTransition11;
            Q[1, 0] = 1.0 - PTransition22;
            Q[1, 1] = PTransition22;
            return Q;
        }

        public override double LogLikelihood(Vector parameter, double penaltyFactor, bool fillOutputs)
        {
            Vector pbak = Parameters;
            if (parameter != null)
                Parameters = parameter;
            double ll = ComputeLogLikelihood(Parameters, null, null, null, null, null);
            if (fillOutputs)
                ComputeResidualsAndOutputs();
            if (parameter != null)
                Parameters = pbak;
            return ll;
        }

        private double ComputeLogLikelihood(Vector parameter, TimeSeries resids, TimeSeries predictiveP1,
                                            TimeSeries predictiveMeans, TimeSeries MRIndicator, TimeSeries TFIndicator)
        {
            Vector pbak = Parameters;
            Parameters = parameter;
            int n = values.Count;
            var pstate = new Vector(new[] {MarginalP1, MarginalP2});  // predictive d-n of state
            var Q = BuildTransitionMatrix();
            Q.Transpose();

            double logLike = 0;

            for (int t=0 ; t<n ; ++t )
            {
                double previousY = t > 0 ? (values[t - 1] - Mu) : (values[0] - Mu);
                double mu1 = Phi1*previousY + Mu;
                double mu2 = Phi2*previousY + Mu;
                double lp1 = NormalDistributionExtensions.LogProbabilityDensity(values[t] - mu1, 0, Sigma);
                double lp2 = NormalDistributionExtensions.LogProbabilityDensity(values[t] - mu2, 0, Sigma);

                if (resids != null)
                {
                    double pMean = pstate[0]*mu1 + pstate[1]*mu2;
                    double resid = (values[t] - pMean)/Sigma;
                    resids.Add(values.TimeStamp(t), resid, false);
                }

                // include adjustment to prevent probability under/over-flows
                double adjust = Math.Max(lp1, lp2);
                lp1 -= adjust;
                lp2 -= adjust;

                // update with log-likelihood of this value
                double mp = pstate[0]*Math.Exp(lp1) + pstate[1]*Math.Exp(lp2); // this is real prob. divided by exp(adjust)
                logLike += Math.Log(mp) + adjust;

                // update to get filtering d-n of state
                var fstate = new Vector(2);
                fstate[0] = pstate[0]*Math.Exp(lp1);
                fstate[1] = pstate[1]*Math.Exp(lp2);
                fstate *= 1.0/(fstate[0] + fstate[1]);

                pstate = (Q*fstate.ToColumnMatrix()).ToVector();           // predictive for next time 
                pstate = pstate*1.0/(pstate[0]+pstate[1]);                 // this prevents small cumulative errors

                double thisY = values[t] - Mu;
                double thisMu1 = Phi1 * thisY + Mu;
                double thisMu2 = Phi2*thisY + Mu;

                if (predictiveMeans != null)
                    predictiveMeans.Add(values.TimeStamp(t), pstate[0]*thisMu1 + pstate[1]*thisMu2, false);

                if (predictiveP1 != null)
                    predictiveP1.Add(values.TimeStamp(t), pstate[0] - MarginalP1, false);

                double mrindic = (pstate[0] - MarginalP1) > 0 ? (Math.Sign(Phi1*values[t]) * (pstate[0]-MarginalP1)) : 0;
                double tfindic = (pstate[0] - MarginalP1) < 0 ? (Math.Sign(Phi2*values[t]) * (MarginalP1-pstate[0])) : 0;
                if (MRIndicator != null)
                    MRIndicator.Add(values.TimeStamp(t), mrindic, false);
                if (TFIndicator != null)
                    TFIndicator.Add(values.TimeStamp(t), tfindic, false);
            }

            Parameters = pbak;
            return logLike;
        }

        protected override Vector ComputeConsequentialParameters(Vector parameter)
        {
            Vector pbak = Parameters;

            Parameters = parameter;
            if (ParameterStates[0] == ParameterState.Consequential)
                Mu = values.SampleMean();
            if (ParameterStates[1] == ParameterState.Consequential)
            {
                double v = values.SampleVariance();
                double s2 = v*(1 - MarginalP1*Phi1*Phi1 -MarginalP2*Phi2*Phi2);
                Sigma = Math.Sqrt(s2);
            }
            if (ParameterStates[2] == ParameterState.Consequential)
                throw new ApplicationException("Phi1 cannot be consequential.");
            if (ParameterStates[3] == ParameterState.Consequential)
                throw new ApplicationException("Phi2 cannot be consequential.");
            if (ParameterStates[4] == ParameterState.Consequential)
                throw new ApplicationException("P11 cannot be consequential.");
            if (ParameterStates[5] == ParameterState.Consequential)
                throw new ApplicationException("P22 cannot be consequential.");

            Vector retval = Parameters;

            Parameters = pbak;
            return retval;
        }

        public override object SimulateData(object inputs, int randomSeed)
        {
            var times = inputs as List<DateTime>;
            if (times == null)
                return null; // inputs should be a list of DateTimes

            int n = times.Count;
            var simulated = new TimeSeries();

            var randomSource = new AdditiveLaggedFibonacciRandomSource(randomSeed);
            var stdnormal = new StandardDistribution(randomSource);
            var cud = new ContinuousUniformDistribution(randomSource);

            int curState = cud.NextDouble() < MarginalP1 ? 1 : 2;
            double previousValue = Mu;

            for (int t=0 ; t<n ; ++t)
            {
                double next;
                if (curState == 1)
                    next = (previousValue - Mu)*Phi1;
                else
                    next = (previousValue - Mu)*Phi2;
                next += Mu + stdnormal.NextDouble()*Sigma;
                simulated.Add(times[t], next, false);
                previousValue = next;

                double pStateChange = curState == 1 ? 1 - PTransition11 : 1 - PTransition22;
                if (cud.NextDouble() <= pStateChange)
                    curState = 3 - curState;
            }

            simulated.Title = "Simul.";
            simulated.Description = "Simulation from " + Description;
            return simulated;
        }

        public override object BuildForecasts(object otherData, object inputs)
        {
            MessageBox.Show("Forecasting not yet implemented for Switching AR models.", "Error");
            return null;
        }

        private void LocalInitializeParameters()
        {
            Parameters = new Vector(6);
            Mu = 0;
            Sigma = 1;
            Phi1 = 0;
            Phi2 = 0;
            PTransition11 = 0.95;
            PTransition22 = 0.95;

            ParameterStates = new ParameterState[Parameters.Length];
            ParameterStates[0] = ParameterState.Consequential; // mu follows
            ParameterStates[1] = ParameterState.Consequential; // sigma follows from the others
            for (int i = 2; i < Parameters.Length; ++i)
                ParameterStates[i] = ParameterState.Free;      // the 2 ar coefficients are free
        }


        protected override void InitializeParameters()
        {
            LocalInitializeParameters();
        }

        public override Vector ComputeACF(int maxLag, bool normalize)
        {
            // garbage for now!  (WN acvf)
            var retval = new Vector(maxLag+1);
            retval[0] = 1.0;
            return retval;
        }

        public override int NumOutputs()
        {
            return 6;
        }

        public override string GetOutputName(int index)
        {
            if (index < 2)
                return base.GetOutputName(index);
            switch (index)
            {
                case 2:
                    return "Predictive P(MR) for Next on Avail.";
                case 3:
                    return "Predictive Mean on Avail.";
                case 4:
                    return "Leading MR Indicator";
                case 5:
                    return "Leading TF Indicator";
            }
            return null;
        }

        public override object GetOutput(int socket)
        {
            if (socket < 2)
                return base.GetOutput(socket);
            switch (socket)
            {
                case 2:
                    return predictiveP1;
                case 3:
                    return predictiveMean;
                case 4:
                    return MRIndicator;
                case 5:
                    return TFIndicator;
            }
            return null;
        }

        #region Real-Time Prediction Functions

        [NonSerialized]
        private Vector predictiveStateDn;
        [NonSerialized]
        private double previousValue;
        [NonSerialized]
        private DistributionSummary predObsDn;

        public double TrendIndicator { get; protected set; }
        public double MeanReversionIndicator { get; protected set; }

        public void ResetRealTimePrediction()
        {
            predictiveStateDn = new Vector(new[] { 0.5, 0.5 });
            previousValue = Mu;
            predObsDn = new DistributionSummary {Mean = 0, StdDev = 1};
        }

        public void Register(TimeSeries series)
        {
            for (int t = 0; t < series.Count; ++t )
                Register(series.TimeStamp(t), series[t]);
        }

        public double Register(DateTime timeStamp, double value, double[] auxiliaryValues)
        {
            return Register(timeStamp, value); // aux. value is irrelevant
        }

        public double Register(DateTime timeStamp, double value)
        {
            var Q = BuildTransitionMatrix();
            Q.Transpose();
            
            var nd = new NormalDistribution(0, Sigma);
            double mu1 = Phi1 * (previousValue-Mu) + Mu;
            double mu2 = Phi2 * (previousValue-Mu) + Mu;
            double p1 = nd.ProbabilityDensity(value - mu1);
            double p2 = nd.ProbabilityDensity(value - mu2);


            // update to get filtering d-n of state

            var fstate = new Vector(2);
            fstate[0] = predictiveStateDn[0] * p1;
            fstate[1] = predictiveStateDn[1] * p2;
            fstate *= 1.0/(fstate[0] + fstate[1]);

            predictiveStateDn = (Q * fstate.ToColumnMatrix()).ToVector();  // predictive for next time 

            MeanReversionIndicator = (predictiveStateDn[0] - MarginalP1) > 0 ? (Math.Sign(Phi1 * value) * (predictiveStateDn[0] - MarginalP1)) : 0;
            TrendIndicator = (predictiveStateDn[0] - MarginalP1) < 0 ? (Math.Sign(Phi2 * value) * (MarginalP1 - predictiveStateDn[0])) : 0;

            double aux = predictiveStateDn[0];

            // now build predictor
            var component1 = new DistributionSummary();
            var component2 = new DistributionSummary();
            component1.Mean = Phi1 * (value - Mu) + Mu;
            component1.StdDev = Sigma;
            component2.Mean = Phi2 * (value - Mu) + Mu;
            component2.StdDev = Sigma;

            //component1.FillGaussianQuantiles(0.005);
            //component2.FillGaussianQuantiles(0.005);
            //predObsDn = component1.MixWith(component2, predictiveStateDn[0], 400);

            previousValue = value;

            return aux;
        }


        public DistributionSummary GetCurrentPredictor(DateTime futureTime)
        {
            return predObsDn;
        }

        #endregion
    }
}